import lazyllm
from lazyllm import pipeline, warp, bind
from lazyllm.components.formatter import JsonFormatter

toc_prompt=""" 你现在是一个智能助手。你的任务是理解用户的输入，将大纲以列表嵌套字典的列表。每个字典包含一个 `title` 和 `describe`，其中 `title` 中需要用Markdown格式标清层级，`describe` 是对该段的描述和写作指导。

请根据以下用户输入生成相应的列表嵌套字典：

输出示例:
[
    {
        "title": "# 一级标题",
        "describe": "请详细描述此标题的内容，提供背景信息和核心观点。"
    },
    {
        "title": "## 二级标题",
        "describe": "请详细描述标题的内容，提供具体的细节和例子来支持一级标题的观点。"
    },
    {
        "title": "### 三级标题",
        "describe": "请详细描述标题的内容，深入分析并提供更多的细节和数据支持。"
    }
]
用户输入如下：
"""

completion_prompt="""
你现在是一个智能助手。你的任务是接收一个包含 `title` 和 `describe` 的字典，并根据 `describe` 中的指导展开写作
输入示例:
{
    "title": "# 一级标题",
    "describe": "这是写作的描述。"
}

输出:
这是展开写作写的内容
接收如下：

"""

writer_prompt = {"system": completion_prompt, "user": '{"title": {title}, "describe": {describe}}'}
with pipeline() as ppl:
    ppl.outline_writer = lazyllm.OnlineChatModule(source='sensenova', model='DeepSeek-V3-1').formatter(JsonFormatter()).prompt(toc_prompt)
    # 规范化：如果大纲返回为单个 dict，则包装成包含一个元素的 list，确保后续 Warp 可正确拆分
    ppl.outline_normalizer = (lambda outlines: outlines if isinstance(outlines, (list, tuple)) else [outlines])
    # 基于规范化后的大纲逐条生成内容
    ppl.story_generater = warp(ppl.outline_writer.share(prompt=writer_prompt).formatter())
    # 合成时也使用规范化后的大纲，避免将 dict 误当作可迭代键集合
    ppl.synthesizer = (lambda *storys, outlines: "\n".join([f"{o['title']}\n{s}" for s, o in zip(storys, outlines)])) | bind(outlines=ppl.output('outline_normalizer'))
lazyllm.WebModule(ppl, port=23466).start().wait()